# from typing_extensions import ParamSpec
from typing import KeysView

from scipy.stats.stats import brunnermunzel
from graph_cut import get_hist, get_histx, get_mask, graph_cut, mask_add_depth

import argparse
# import torch
# import torch.nn as nn
# from torch.utils.data import DataLoader
from tqdm import tqdm
import numpy as np
import math
import os
import json
import cv2
# torch.backends.cudnn.enabled = False
import matplotlib.pyplot as plt


def get_geometry_feat(mask):
    
    mask_raw = mask*255
    # cv2.cvtColor
    #轮廓
    contours,hierarchy = cv2.findContours(mask_raw,cv2.RETR_LIST,cv2.CHAIN_APPROX_SIMPLE)
    area = []
    for i in range(len(contours)):
        area.append(cv2.contourArea(contours[i]))
    print(len(area))
    max_idx = np.argmax(area)
    mask_raw = np.zeros_like(mask_raw)
    cv2.drawContours(mask_raw, contours, max_idx, 1, 0)
    perimeter = np.sum(mask_raw)
    cv2.fillPoly(mask_raw, [contours[max_idx]], 1)
    area = np.sum(mask_raw)
    return perimeter, area

def get_color_feat(rgb_result):
    img = cv2.cvtColor(rgb_result, cv2.COLOR_BGR2HSV)
    img = np.swapaxes(np.swapaxes(img, 1, 2), 0, 1)

    cnt = np.bincount(np.ndarray.flatten(img[0]))
    feat_h = np.argmax(cnt[1:])+1
    cnt = np.bincount(np.ndarray.flatten(img[1]))
    feat_s = np.argmax(cnt[1:])+1

    return feat_h, feat_s

def get_depth_feat(depth_result):
    img = depth_result
    cnt = np.bincount(np.ndarray.flatten(img))
    feat_d = np.argmax(cnt[1:])+1

    return feat_d
def get_roi(img):
    img = img[:, 420:1500]
    img = cv2.resize(img, (650, 650))
    img = img[20:20+608, 20:20+608]
    return img
def calc_2D_Entropy(img):
    '''
    邻域 3*3的小格子
     __ __ __
    |__|__|__|
    |__||||__|
    |__|__|__|
    角点
     __ __
    ||||__|
    |__|__|
    边
     __ __
    |  |__|
    ||||__|
    |__|__|
    '''
    img = get_roi(img)
    N = 1
    S=img.shape
    # print(S)
    IJ = []
    #计算j
    for row in range(S[0]):
        for col in range(S[1]):
            Left_x=np.max([0,col-N])
            Right_x=np.min([S[1],col+N+1])
            up_y=np.max([0,row-N])
            down_y=np.min([S[0],row+N+1])
            region=img[up_y:down_y,Left_x:Right_x] # 九宫格区域
            j = (np.sum(region) - img[row][col])/((2*N+1)**2-1)
            IJ.append([img[row][col],j])
    # print(IJ)
    # 计算F(i,j)
    F=[]
    arr = [list(i) for i in set(tuple(j) for j in IJ)] #去重，会改变顺序，不过此处不影响
    for i in range(len(arr)):
        F.append(IJ.count(arr[i]))
    # print(F)
    # 计算pij
    P=np.array(F)/len(F)

    # 计算熵
    
    E = np.sum([p *np.log2(1/p) for p in P])
    # print(E)
    return E

def get_color_entropy(rgb_result):
    img = cv2.cvtColor(rgb_result, cv2.COLOR_BGR2HSV)
    img = np.swapaxes(np.swapaxes(img, 1, 2), 0, 1)

    entropy_h = calc_2D_Entropy(img[0])
    entropy_s = calc_2D_Entropy(img[1])

    return entropy_h, entropy_s
def get_depth_entropy(depth_result):
    img = depth_result

    entropy_d = calc_2D_Entropy(img)
    return entropy_d

parser = argparse.ArgumentParser(description='Greenhouse Image Analysis')
parser.add_argument('--image_dir',   type=str,   default='C:\\Users\\ys\\Desktop\\data_original\\', help='image_dir')
parser.add_argument('--gt_file',     type=str,   default='GroundTruth_All_388_Images.json', help='gt_file')
# parser.add_argument('--output_dir',   type=str,   default='./output_whitebg/', help='output_dir')
# parser.add_argument('--variety',   type=str,   default='Salanova', help='Aphylion, Salanova, Satine, Lugano')
parser.add_argument('--startn',   type=int,   default='250', help='Aphylion, Salanova, Satine, Lugano')
params = parser.parse_args()
print(params)

json_path = os.path.join(params.image_dir, f'GroundTruth_All_388_Images.json')
jsn_f = open(json_path, 'r')
jsn = json.load(jsn_f)
jsn_dataset = jsn['Measurements']

a_size = 15

keys = []
i = 0
for key, value in jsn_dataset.items():
    if i >= params.startn and i < params.startn + 50:
        keys.append(key)
    i += 1
json_all = {}

i = 0
for key in tqdm(keys, desc=str(params.startn)):
    i += 1
    key = 'Image4'
    value=jsn_dataset[key]
    output_dir = f'./data_graphcut/'
    # output_dir = f'./output/{value["Variety"]}/'
    rgb_filename = value['RGB_Image']
    depth_filename = value['Depth_Information']
    histx = get_histx(value["Variety"])

    # print(value["Variety"])
    img_rgb_path = os.path.join(params.image_dir+'RGBImages', rgb_filename)
    img_depth_path = os.path.join(params.image_dir+'DepthImages', depth_filename)
    img_rgb = cv2.imread(img_rgb_path)
    img_depth = cv2.imread(img_depth_path, cv2.IMREAD_UNCHANGED)
    img_depth = np.where(img_depth == 0, 65535, img_depth)

    
    mask = get_mask(img_rgb, a_size, histx)
    mask = mask_add_depth(img_depth, mask)
    # cv2.imwrite(output_dir+rgb_filename[:-4]+'.jpg', mask*70)

    if i > 12:
        print(img_rgb_path)
        print(img_depth_path)
        cv2.imshow('img_rgb_before', img_rgb)
        # cv2.imshow('rgb_result', rgb_result)
        
        
        cv2.imshow('mask_before', mask*70)
        # cv2.waitKey(0)
    rgb_result, depth_result, mask = graph_cut(img_rgb, img_depth, mask=mask)
    if i > 12:
        print(img_rgb_path)
        cv2.imshow('img_rgb_after', img_rgb)
        # cv2.imshow('rgb_result', rgb_result)
        
        
        cv2.imshow('mask_after', mask*70)
        cv2.waitKey(0)
    print(key)

    perimeter, area = get_geometry_feat(mask)
    feat_h, feat_s = get_color_feat(rgb_result)
    feat_d = get_depth_feat(depth_result)
    value['Perimeter'] = float(perimeter)
    value['Area'] = float(area)
    value['Feat_h'] = float(feat_h)
    value['Feat_s'] = float(feat_s)
    value['Feat_d'] = float(feat_d)
    
    # entropy_h, entropy_s = get_color_entropy(rgb_result)
    # value['Entropy_h'] = float(entropy_h)
    # value['Entropy_s'] = float(entropy_s)
    
    json_all[key] = value
    # json_all.update(jsn_dataset)
    results = {'Measurements': json_all}
    new_jsn = json.dumps(results)
    new_jsn_file = open(os.path.join(output_dir, 'feat'+str(params.startn)+'.json'), 'w+')
    new_jsn_file.write(new_jsn)



    # os.makedirs(os.path.join(output_dir+'RGBImages'), exist_ok=True)
    # os.makedirs(os.path.join(output_dir+'DepthImages'), exist_ok=True)
    # cv2.imwrite(os.path.join(output_dir, 'RGBImages', rgb_filename), rgb_result)
    # cv2.imwrite(os.path.join(output_dir, 'DepthImages', depth_filename), depth_result)
    
    # r = (486, 914, 486+155+50, 914+167+100)
    # img_ori = cv2.rectangle(img_ori, (r[1], r[0]), (r[3], r[2]), (0, 0, 0), 3)
    # cv2.imshow('img', depth_result*200)
    # cv2.waitKey(300)

# cv2.imwrite(f'./output/{filename}', result)


